{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "import dgl\n",
    "from dgl import DGLGraph\n",
    "from dgl.data import MiniGCDataset\n",
    "import dgl.function as fn\n",
    "\n",
    "from functools import partial\n",
    "\n",
    "import re\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from pytorch_pretrained_bert import BertTokenizer\n",
    "\n",
    "import spacy\n",
    "import pickle\n",
    "import collections\n",
    "\n",
    "from sklearn.metrics import log_loss\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'\n",
    "torch.cuda.set_device(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "Load results from BERT Embeddings Generatation with Punc.ipynb and BERT Embeddings Generatation without Punc.ipynb\n",
    "'''\n",
    "\n",
    "token_lst = pickle.load(open('token_lst_wto_padding.pkl', \"rb\")) # tokens of every sentence without padding\n",
    "bert_outputs_lst = pickle.load(open('bert_outputs.pkl', \"rb\")) # list of outputs of bert for every sentence\n",
    "offsets_lst = pickle.load(open('offsets_lst.pkl', \"rb\"))\n",
    "\n",
    "test_token_lst = pickle.load(open('test_token_lst_wto_padding.pkl', \"rb\")) # tokens of every sentence without padding\n",
    "test_bert_outputs_lst = pickle.load(open('test_bert_outputs.pkl', \"rb\")) # list of outputs of bert for every sentence\n",
    "test_offsets_lst = pickle.load(open('test_offsets_lst.pkl', \"rb\"))\n",
    "\n",
    "others_bert_outputs = pickle.load(open('others_bert_outputs.pkl', \"rb\"))\n",
    "test_others_bert_outputs  = pickle.load(open('test_others_bert_outputs.pkl', \"rb\"))\n",
    "\n",
    "\n",
    "train_df = pd.concat([\n",
    "    pd.read_csv(\"gap-test.tsv\", delimiter=\"\\t\"),\n",
    "    pd.read_csv(\"gap-validation.tsv\", delimiter=\"\\t\")\n",
    "], axis=0)\n",
    "\n",
    "test_df = pd.read_csv(\"gap-development.tsv\", delimiter=\"\\t\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Building"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RGCNLayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RGCNLayer(nn.Module):\n",
    "    def __init__(self, feat_size, num_rels, activation=None, gated = True):\n",
    "        \n",
    "        super(RGCNLayer, self).__init__()\n",
    "        self.feat_size = feat_size\n",
    "        self.num_rels = num_rels\n",
    "        self.activation = activation\n",
    "        self.gated = gated\n",
    "\n",
    "        self.weight = nn.Parameter(torch.Tensor(self.num_rels, self.feat_size, 256))\n",
    "        # init trainable parameters\n",
    "        nn.init.xavier_uniform_(self.weight,gain=nn.init.calculate_gain('relu'))\n",
    "        \n",
    "        if self.gated:\n",
    "            self.gate_weight = nn.Parameter(torch.Tensor(self.num_rels, self.feat_size, 1))\n",
    "            nn.init.xavier_uniform_(self.gate_weight,gain=nn.init.calculate_gain('sigmoid'))\n",
    "        \n",
    "    def forward(self, g):\n",
    "        \n",
    "        weight = self.weight\n",
    "        gate_weight = self.gate_weight\n",
    "        \n",
    "        def message_func(edges):\n",
    "            w = weight[edges.data['rel_type']]\n",
    "            msg = torch.bmm(edges.src['h'].unsqueeze(1), w).squeeze()\n",
    "            msg = msg * edges.data['norm']\n",
    "            \n",
    "            if self.gated:\n",
    "                gate_w = gate_weight[edges.data['rel_type']]\n",
    "                gate = torch.bmm(edges.src['h'].unsqueeze(1), gate_w).squeeze().reshape(-1,1)\n",
    "                gate = torch.sigmoid(gate)\n",
    "                msg = msg * gate\n",
    "                \n",
    "            return {'msg': msg}\n",
    "    \n",
    "        def apply_func(nodes):\n",
    "            h = nodes.data['h']\n",
    "            h = self.activation(h)\n",
    "            return {'h': h}\n",
    "\n",
    "        g.update_all(message_func, fn.sum(msg='msg', out='h'), apply_func)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Full RGCN Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RGCNModel(nn.Module):\n",
    "    def __init__(self, h_dim, num_rels, num_hidden_layers=1, gated = True):\n",
    "        super(RGCNModel, self).__init__()\n",
    "\n",
    "        self.h_dim = h_dim\n",
    "        self.num_rels = num_rels\n",
    "        self.num_hidden_layers = num_hidden_layers\n",
    "        self.gated = gated\n",
    "        \n",
    "        # create rgcn layers\n",
    "        self.build_model()\n",
    "       \n",
    "    def build_model(self):        \n",
    "        self.layers = nn.ModuleList() \n",
    "        for _ in range(self.num_hidden_layers):\n",
    "            rgcn_layer = RGCNLayer(self.h_dim, self.num_rels, activation=F.relu, gated = self.gated)\n",
    "            self.layers.append(rgcn_layer)\n",
    "    \n",
    "    def forward(self, g):\n",
    "        for layer in self.layers:\n",
    "            layer(g)\n",
    "        \n",
    "        rst_hidden = []\n",
    "        for sub_g in dgl.unbatch(g):\n",
    "            rst_hidden.append(  sub_g.ndata['h']   )\n",
    "        return rst_hidden"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Design the Main Model (R-GCN + FFNN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Head(nn.Module):\n",
    "    \"\"\"The MLP submodule\"\"\"\n",
    "    def __init__(self, gcn_out_size: int, bert_out_size: int):\n",
    "        super().__init__()\n",
    "        self.bert_out_size = bert_out_size\n",
    "        self.gcn_out_size = gcn_out_size\n",
    "        \n",
    "        self.fc = nn.Sequential(\n",
    "            nn.BatchNorm1d(bert_out_size * 3 + gcn_out_size * 3),\n",
    "            nn.Dropout(0.5),\n",
    "            nn.Linear(bert_out_size * 3 + gcn_out_size * 3, 256),    \n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm1d(256),\n",
    "            nn.Dropout(0.5),\n",
    "            nn.Linear(256, 3),\n",
    "        )\n",
    "        for i, module in enumerate(self.fc):\n",
    "            if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d)):\n",
    "                nn.init.constant_(module.weight, 1)\n",
    "                nn.init.constant_(module.bias, 0)\n",
    "            elif isinstance(module, nn.Linear):\n",
    "                if getattr(module, \"weight_v\", None) is not None:\n",
    "                    nn.init.uniform_(module.weight_g, 0, 1)\n",
    "                    nn.init.kaiming_normal_(module.weight_v)\n",
    "                    assert model[i].weight_g is not None\n",
    "                else:\n",
    "                    nn.init.kaiming_normal_(module.weight)\n",
    "                nn.init.constant_(module.bias, 0)\n",
    "                \n",
    "    def forward(self, gcn_outputs, offsets_gcn, bert_embeddings):\n",
    "        \n",
    "        gcn_extracted_outputs = [gcn_outputs[i].unsqueeze(0).gather(1, offsets_gcn[i].unsqueeze(0).unsqueeze(2)\n",
    "                                       .expand(-1, -1, gcn_outputs[i].unsqueeze(0).size(2))).view(gcn_outputs[i].unsqueeze(0).size(0), -1) for i in range(len(gcn_outputs))]\n",
    "        \n",
    "        gcn_extracted_outputs = torch.stack(gcn_extracted_outputs, dim=0).squeeze()\n",
    "        \n",
    "        embeddings = torch.cat((gcn_extracted_outputs, bert_embeddings), 1) \n",
    "        \n",
    "        return self.fc(embeddings)\n",
    "\n",
    "\n",
    "class BERT_Head(nn.Module):\n",
    "    def __init__(self, bert_hidden_size: int):\n",
    "        super().__init__()\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.BatchNorm1d(bert_hidden_size * 3),\n",
    "            nn.Dropout(0.5),\n",
    "            nn.Linear(bert_hidden_size * 3, 512 * 3),   \n",
    "            nn.ReLU(),\n",
    "        )\n",
    "\n",
    "        for i, module in enumerate(self.fc):\n",
    "            if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d)):\n",
    "                nn.init.constant_(module.weight, 1)\n",
    "                nn.init.constant_(module.bias, 0)\n",
    "            elif isinstance(module, nn.Linear):\n",
    "                if getattr(module, \"weight_v\", None) is not None:\n",
    "                    nn.init.uniform_(module.weight_g, 0, 1)\n",
    "                    nn.init.kaiming_normal_(module.weight_v)\n",
    "                    assert model[i].weight_g is not None\n",
    "                else:\n",
    "                    nn.init.kaiming_normal_(module.weight)\n",
    "                nn.init.constant_(module.bias, 0)\n",
    "                \n",
    "    def forward(self, bert_embeddings):\n",
    "        #print('BERT_Head bert_embeddings: ', bert_embeddings, bert_embeddings.view(bert_embeddings.shape[0],-1).shape)\n",
    "        outputs = self.fc(bert_embeddings.view(bert_embeddings.shape[0],-1))\n",
    "        return outputs\n",
    "        \n",
    "        \n",
    "    \n",
    "    \n",
    "    \n",
    "class GPRModel(nn.Module):\n",
    "    \"\"\"The main model.\"\"\"\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.RGCN =  RGCNModel(h_dim = 1024, num_rels = 3, gated = True)\n",
    "        self.BERThead = BERT_Head(1024) # bert output size\n",
    "        self.head = Head(256, 512)  # gcn output   berthead output\n",
    "    \n",
    "    \n",
    "    def forward(self, offsets_bert, offsets_gcn, bert_embeddings, g):\n",
    "        gcn_outputs = self.RGCN(g)\n",
    "        bert_head_outputs = self.BERThead(bert_embeddings)\n",
    "        head_outputs = self.head(gcn_outputs, offsets_gcn, bert_head_outputs)\n",
    "        return head_outputs            \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Input"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate All Syntactic Graphs with BERT embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = spacy.load('en_core_web_lg')\n",
    "\n",
    "BERT_MODEL = 'bert-large-uncased'\n",
    "tokenizer = BertTokenizer.from_pretrained(BERT_MODEL, never_split = (\"[UNK]\", \"[SEP]\", \"[PAD]\", \"[CLS]\", \"[MASK]\", \"[THISISA]\", \"[THISISB]\", \"[THISISP]\"))\n",
    "\n",
    "tokenizer.vocab[\"[THISISA]\"] = -1\n",
    "tokenizer.vocab[\"[THISISB]\"] = -1\n",
    "tokenizer.vocab[\"[THISISP]\"] = -1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_target(i, target_offset_list):\n",
    "    return i in target_offset_list\n",
    "\n",
    "def transfer_n_e(nodes, edges):\n",
    "\n",
    "    num_nodes = len(nodes)\n",
    "    new_edges = []\n",
    "    for e1, e2 in edges:\n",
    "        new_edges.append( [nodes[e1], nodes[e2]] ) \n",
    "    return num_nodes, new_edges\n",
    "\n",
    "all_graphs = []\n",
    "gcn_offsets = []\n",
    "for i, sent_token in enumerate(token_lst):\n",
    "    sent_token = token_lst[i]\n",
    "\n",
    "    sent = ' '.join([re.sub(\"[#]\",\"\",token)   for token in tokenizer.convert_ids_to_tokens(sent_token[1:-1])])\n",
    "    doc = parser(sent)\n",
    "    parse_rst = doc.to_json()\n",
    "\n",
    "    target_offset_list = [item - 1 for item in offsets_lst[i]]\n",
    "    \n",
    "    nodes = collections.OrderedDict()\n",
    "    edges = []\n",
    "    edge_type = []\n",
    "    for i_word, word in enumerate(parse_rst['tokens']):\n",
    "        if not (is_target(i_word, target_offset_list) or is_target(word['head'], target_offset_list)):\n",
    "            continue\n",
    "\n",
    "        if i_word not in nodes:\n",
    "            nodes[i_word] = len(nodes) \n",
    "            edges.append( [i_word, i_word] )\n",
    "            edge_type.append(0)\n",
    "        if word['head'] not in nodes:\n",
    "            nodes[word['head']] = len(nodes) \n",
    "            edges.append( [word['head'], word['head']] )\n",
    "            edge_type.append(0)\n",
    "\n",
    "        if word['dep'] != 'ROOT':\n",
    "                edges.append( [word['head'], word['id']] )\n",
    "                edge_type.append(1)\n",
    "                edges.append( [word['id'], word['head']] )\n",
    "                edge_type.append(2)\n",
    "\n",
    "    num_nodes, tran_edges = transfer_n_e(nodes, edges)\n",
    "    \n",
    "    gcn_offset = [nodes[offset] for offset in target_offset_list]\n",
    "    gcn_offsets.append(gcn_offset)\n",
    "    \n",
    "    G = dgl.DGLGraph()\n",
    "    G.add_nodes(num_nodes)\n",
    "    G.add_edges(list(zip(*tran_edges))[0],list(zip(*tran_edges))[1]) \n",
    "\n",
    "    for i_word, word in enumerate(parse_rst['tokens']):\n",
    "        if not (is_target(i_word, target_offset_list) or is_target(word['head'], target_offset_list)):\n",
    "            continue\n",
    "        if is_target(i_word, target_offset_list): \n",
    "            G.nodes[ [ nodes[i_word] ]].data['h'] = others_bert_outputs[i][0][target_offset_list.index(i_word)].unsqueeze(0).cuda()\n",
    "        else:\n",
    "            G.nodes[ [ nodes[i_word] ]].data['h'] = bert_outputs_lst[i][0][i_word + 1].unsqueeze(0).cuda()\n",
    "        if is_target(word['head'], target_offset_list):\n",
    "            G.nodes[ [ nodes[word['head']] ]].data['h'] = others_bert_outputs[i][0][target_offset_list.index(word['head'])].unsqueeze(0).cuda()\n",
    "        else:   \n",
    "            G.nodes[ [ nodes[word['head']] ]].data['h'] = bert_outputs_lst[i][0][word['head'] + 1].unsqueeze(0).cuda()\n",
    "\n",
    "    edge_norm = []\n",
    "    for e1, e2 in tran_edges:\n",
    "        if e1 == e2:\n",
    "            edge_norm.append(1)\n",
    "        else:\n",
    "            edge_norm.append( 1 / (G.in_degree(e2) - 1 ) )\n",
    "\n",
    "\n",
    "    edge_type = torch.from_numpy(np.array(edge_type))\n",
    "    edge_norm = torch.from_numpy(np.array(edge_norm)).unsqueeze(1).float().cuda()\n",
    "\n",
    "    G.edata.update({'rel_type': edge_type,})\n",
    "    G.edata.update({'norm': edge_norm})\n",
    "    all_graphs.append(G)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_all_graphs = []\n",
    "test_gcn_offsets = []\n",
    "for i, sent_token in enumerate(test_token_lst):\n",
    "    sent_token = test_token_lst[i]\n",
    "\n",
    "    sent = ' '.join([re.sub(\"[#]\",\"\",token)   for token in tokenizer.convert_ids_to_tokens(sent_token[1:-1])])\n",
    "    doc = parser(sent)\n",
    "    parse_rst = doc.to_json()\n",
    "\n",
    "    target_offset_list = [item - 1 for item in test_offsets_lst[i]]\n",
    "    \n",
    "    nodes = collections.OrderedDict()\n",
    "    edges = []\n",
    "    edge_type = []\n",
    "    for i_word, word in enumerate(parse_rst['tokens']):\n",
    "        if not (is_target(i_word, target_offset_list) or is_target(word['head'], target_offset_list)):\n",
    "            continue\n",
    "\n",
    "        if i_word not in nodes:\n",
    "            nodes[i_word] = len(nodes) \n",
    "            edges.append( [i_word, i_word] )\n",
    "            edge_type.append(0)\n",
    "        if word['head'] not in nodes:\n",
    "            nodes[word['head']] = len(nodes) \n",
    "            edges.append( [word['head'], word['head']] )\n",
    "            edge_type.append(0)\n",
    "\n",
    "        if word['dep'] != 'ROOT':\n",
    "                edges.append( [word['head'], word['id']] )\n",
    "                edge_type.append(1)\n",
    "                edges.append( [word['id'], word['head']] )\n",
    "                edge_type.append(2)\n",
    "\n",
    "    num_nodes, tran_edges = transfer_n_e(nodes, edges)\n",
    "    \n",
    "    test_gcn_offset = [nodes[offset] for offset in target_offset_list]\n",
    "    test_gcn_offsets.append(test_gcn_offset)\n",
    "    \n",
    "    G = dgl.DGLGraph()\n",
    "    G.add_nodes(num_nodes)\n",
    "    G.add_edges(list(zip(*tran_edges))[0],list(zip(*tran_edges))[1]) \n",
    "\n",
    "    for i_word, word in enumerate(parse_rst['tokens']):\n",
    "        if not (is_target(i_word, target_offset_list) or is_target(word['head'], target_offset_list)):\n",
    "            continue\n",
    "        \n",
    "        if is_target(i_word, target_offset_list): \n",
    "            G.nodes[ [ nodes[i_word] ]].data['h'] = test_others_bert_outputs[i][0][target_offset_list.index(i_word)].unsqueeze(0).cuda()\n",
    "        else:\n",
    "            G.nodes[ [ nodes[i_word] ]].data['h'] = test_bert_outputs_lst[i][0][i_word + 1].unsqueeze(0).cuda()\n",
    "        if is_target(word['head'], target_offset_list):\n",
    "            G.nodes[ [ nodes[word['head']] ]].data['h'] = test_others_bert_outputs[i][0][target_offset_list.index(word['head'])].unsqueeze(0).cuda()\n",
    "        else:   \n",
    "            G.nodes[ [ nodes[word['head']] ]].data['h'] = test_bert_outputs_lst[i][0][word['head'] + 1].unsqueeze(0).cuda()\n",
    "\n",
    "    edge_norm = []\n",
    "    for e1, e2 in tran_edges:\n",
    "        if e1 == e2:\n",
    "            edge_norm.append(1)\n",
    "        else:\n",
    "            edge_norm.append( 1 / (G.in_degree(e2) - 1 ) )\n",
    "\n",
    "\n",
    "    edge_type = torch.from_numpy(np.array(edge_type))\n",
    "    edge_norm = torch.from_numpy(np.array(edge_norm)).unsqueeze(1).float().cuda()\n",
    "\n",
    "    G.edata.update({'rel_type': edge_type,})\n",
    "    G.edata.update({'norm': edge_norm})\n",
    "    test_all_graphs.append(G)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Design Dataloader and Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPRDataset(Dataset):\n",
    "    def __init__(self, original_df, graphs, bert_offsets, gcn_offsets, bert_embeddings):\n",
    "        \n",
    "\n",
    "\n",
    "        tmp = original_df[[\"A-coref\", \"B-coref\"]].copy()\n",
    "        tmp[\"Neither\"] = ~(original_df[\"A-coref\"] | original_df[\"B-coref\"])\n",
    "        self.y = tmp.values.astype(\"bool\")\n",
    "\n",
    "        self.graphs = graphs\n",
    "        self.bert_offsets = bert_offsets  # 已经+1了\n",
    "        self.bert_embeddings = bert_embeddings  # 有[CLS]\n",
    "        self.gcn_offsets = gcn_offsets\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.graphs)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.graphs[idx], self.bert_offsets[idx], self.gcn_offsets[idx], self.bert_embeddings[idx], self.y[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate(samples):\n",
    "    \n",
    "    graphs, bert_offsets, gcn_offsets, bert_embeddings, labels = map(list, zip(*samples))\n",
    "    \n",
    "    batched_graph = dgl.batch(graphs)\n",
    "    offsets_bert = torch.stack([torch.LongTensor(x) for x in bert_offsets], dim=0)\n",
    "    offsets_gcn = torch.stack([torch.LongTensor(x) for x in gcn_offsets], dim=0)\n",
    "    \n",
    "    one_hot_labels = torch.stack([torch.from_numpy(x.astype(\"uint8\")) for x in labels], dim=0)\n",
    "    _, labels = one_hot_labels.max(dim=1)\n",
    "    \n",
    "    bert_embeddings = torch.stack(bert_embeddings, dim=0).squeeze()\n",
    "    \n",
    "    \n",
    "    return batched_graph, offsets_bert, offsets_gcn, bert_embeddings, labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test DataLoarder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset = GPRDataset(original_df = test_df, graphs = test_all_graphs, bert_offsets = test_offsets_lst, gcn_offsets = test_gcn_offsets, bert_embeddings = test_others_bert_outputs)\n",
    "#train_dataset = GPRDataset(original_df = train_df, graphs = all_graphs, bert_offsets = offsets_lst, gcn_offsets= gcn_offsets, bert_embeddings = others_bert_outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#train_dataloarder = DataLoader(\n",
    "#    train_dataset,\n",
    "#    collate_fn = collate,\n",
    "#    batch_size = 4,\n",
    "#    shuffle=True,\n",
    "#)\n",
    "\n",
    "test_dataloarder = DataLoader(\n",
    "    test_dataset,\n",
    "    collate_fn = collate,\n",
    "    batch_size = 4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def send_graph_to_cpu(g):\n",
    "    # nodes\n",
    "    labels = g.node_attr_schemes()\n",
    "    for l in labels.keys():\n",
    "        g.ndata[l] = g.ndata.pop(l).cpu()\n",
    "    # edges\n",
    "    labels = g.edge_attr_schemes()\n",
    "    for l in labels.keys():\n",
    "        g.edata[l] = g.edata.pop(l).cpu()\n",
    "    return g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr_value = 0.0001\n",
    "total_epoch = 100\n",
    "def adjust_learning_rate(optimizers, epoch):\n",
    "    # warm up\n",
    "    if epoch < 10:\n",
    "        lr_tmp = 0.00001\n",
    "    else:\n",
    "        lr_tmp = lr_value * pow((1 - 1.0 * epoch / 100), 0.9)\n",
    "    \n",
    "    if epoch > 36:\n",
    "        lr_tmp =  0.000015 * pow((1 - 1.0 * epoch / 100), 0.9)\n",
    "    \n",
    "    for optimizer in optimizers:\n",
    "        for param_group in optimizer.param_groups:\n",
    "            param_group['lr'] = lr_tmp\n",
    "\n",
    "    return lr_tmp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate labels\n",
    "tmp = train_df[[\"A-coref\", \"B-coref\"]].copy()\n",
    "tmp[\"Neither\"] = ~(train_df[\"A-coref\"] | train_df[\"B-coref\"])\n",
    "train_y = tmp.values.astype(\"bool\").argmax(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5 Fold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from operator import itemgetter\n",
    "\n",
    "kfold = StratifiedKFold(n_splits = 5)\n",
    "test_predict_lst = [] # the test output for every fold\n",
    "for train_index, test_index in kfold.split(train_df, train_y):\n",
    "    print(\"=\" * 20)\n",
    "    print(f\"Fold {len(test_predict_lst) + 1}\")\n",
    "    print(\"=\" * 20)\n",
    "    \n",
    "    val_dataset = GPRDataset(original_df = train_df.iloc[test_index] , \n",
    "                               graphs = list(itemgetter(*test_index)(all_graphs)), \n",
    "                               bert_offsets = list(itemgetter(*test_index)(offsets_lst)), \n",
    "                               gcn_offsets= list(itemgetter(*test_index)(gcn_offsets)) , \n",
    "                               bert_embeddings = list(itemgetter(*test_index)(others_bert_outputs)) )\n",
    "    \n",
    "    train_dataset = GPRDataset(original_df = train_df.iloc[train_index] , \n",
    "                               graphs = list(itemgetter(*train_index)(all_graphs)), \n",
    "                               bert_offsets = list(itemgetter(*train_index)(offsets_lst)), \n",
    "                               gcn_offsets= list(itemgetter(*train_index)(gcn_offsets)) , \n",
    "                               bert_embeddings = list(itemgetter(*train_index)(others_bert_outputs)) )\n",
    "    \n",
    "    train_dataloarder = DataLoader(\n",
    "    train_dataset,\n",
    "    collate_fn = collate,\n",
    "    batch_size = 4,\n",
    "    shuffle=True,)\n",
    "    \n",
    "    val_dataloarder = DataLoader(\n",
    "    val_dataset,\n",
    "    collate_fn = collate,\n",
    "    batch_size = 4,)\n",
    "\n",
    "    model = GPRModel().cuda()\n",
    "    loss_func = nn.CrossEntropyLoss() \n",
    "    optimizer = optim.Adam(model.parameters(), lr=lr_value)\n",
    "    reg_lambda = 0.035\n",
    "\n",
    "    print('Dataloader Success---------------------')\n",
    "    \n",
    "    best_val_loss = 11\n",
    "    #ce_losses = []\n",
    "    #epoch_losses = []\n",
    "    #val_looses = []\n",
    "    for epoch in range(total_epoch):\n",
    "        \n",
    "        if epoch % 5 == 0:\n",
    "            print('|',\">\" * epoch,\" \"*(80-epoch),'|')\n",
    "        \n",
    "        lr = adjust_learning_rate([optimizer],epoch)\n",
    "        #print(\"Learning rate = %4f\\n\" % lr)\n",
    "        model.train()\n",
    "        #epoch_loss = 0\n",
    "        #reg_loss = 0\n",
    "        #ce_loss = 0\n",
    "        for iter, (batched_graph, offsets_bert, offsets_gcn, bert_embeddings, labels) in enumerate(train_dataloarder):\n",
    "\n",
    "            bert_embeddings = bert_embeddings.cuda()\n",
    "            labels = labels.cuda()\n",
    "            offsets_gcn = offsets_gcn.cuda()\n",
    "\n",
    "\n",
    "            prediction = model(offsets_bert, offsets_gcn, bert_embeddings, batched_graph)\n",
    "\n",
    "            l2_reg = None\n",
    "            for w in model.RGCN.parameters():\n",
    "                if not l2_reg:\n",
    "                    l2_reg = w.norm(2)\n",
    "                else:\n",
    "                    l2_reg = l2_reg + w.norm(2)  \n",
    "            for w in model.head.parameters():\n",
    "                if not l2_reg:\n",
    "                    l2_reg = w.norm(2)\n",
    "                else:\n",
    "                    l2_reg = l2_reg + w.norm(2)   \n",
    "            for w in model.BERThead.parameters():\n",
    "                if not l2_reg:\n",
    "                    l2_reg = w.norm(2)\n",
    "                else:\n",
    "                    l2_reg = l2_reg + w.norm(2) \n",
    "            loss = loss_func(prediction, labels) + l2_reg * reg_lambda\n",
    "            #loss = loss_func(prediction, labels)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            #epoch_loss += loss.detach().item()\n",
    "            #reg_loss += (l2_reg*reg_lambda).detach().item()\n",
    "            #ce_loss += (loss_func(prediction, labels)).detach().item()\n",
    "        #epoch_loss /= (iter + 1)\n",
    "        #ce_loss /= (iter + 1)\n",
    "        #reg_loss /= (iter + 1)\n",
    "        #print('Epoch {}, loss {:.4f}, ce_loss {:.4f}, reg_loss {:.4f}'.format(epoch, epoch_loss, ce_loss, reg_loss))\n",
    "        #print('Epoch {}, loss {:.4f}'.format(epoch, epoch_loss))\n",
    "        #epoch_losses.append(epoch_loss)\n",
    "        #ce_losses.append(ce_loss)\n",
    "    \n",
    "        val_loss = 0\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            for iter, (batched_graph, offsets_bert, offsets_gcn, bert_embeddings, labels) in enumerate(val_dataloarder):\n",
    "                offsets_gcn = offsets_gcn.cuda()\n",
    "                bert_embeddings = bert_embeddings.cuda()\n",
    "                labels = labels.cuda()\n",
    "                prediction = model(offsets_bert, offsets_gcn, bert_embeddings, batched_graph)\n",
    "                loss = loss_func(prediction, labels)\n",
    "                val_loss += loss.detach().item()\n",
    "            val_loss = val_loss/(iter + 1)\n",
    "            \n",
    "            \n",
    "        if epoch%20 == 0:\n",
    "            print('Epoch {}, val_loss {:.4f}'.format(epoch, val_loss))\n",
    "\n",
    "        if val_loss < best_val_loss:\n",
    "            best_val_loss = val_loss\n",
    "            if epoch > 20:\n",
    "                torch.save(model.state_dict(), 'best_model.pth') \n",
    "            if epoch > 36: print('Best val loss found: ', best_val_loss)\n",
    "            \n",
    "            \n",
    "            \n",
    "        #print('Epoch {}, val_loss {:.4f}'.format(epoch, val_loss))\n",
    "        #val_looses.append(val_loss)\n",
    "        \n",
    "        #if val_loss < best_val_loss:\n",
    "        #    best_val_loss = val_loss\n",
    "        \n",
    "        #if epoch > 0 and epoch % 5 == 0:\n",
    "        #    plt.figure()\n",
    "        #    plt.plot(ce_losses[len(ce_losses) * 2 // 5 :])\n",
    "        #    plt.plot(val_looses[len(val_looses) * 2 // 5:] )\n",
    "        #    plt.show()\n",
    "    \n",
    "    print('This fold, the best val loss is: ', best_val_loss)\n",
    "    \n",
    "    test_loss = 0.\n",
    "    test_predict = None\n",
    "    \n",
    "    model = GPRModel()\n",
    "    model.load_state_dict(torch.load('best_model.pth', map_location=torch.device('cpu')))\n",
    "    model.eval()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for iter, (batched_graph, offsets_bert, offsets_gcn, bert_embeddings, labels) in enumerate(test_dataloarder):\n",
    "            \n",
    "            offsets_gcn = offsets_gcn\n",
    "            bert_embeddings = bert_embeddings.cpu()\n",
    "            labels = labels\n",
    "            batched_graph = send_graph_to_cpu(batched_graph)\n",
    "            \n",
    "            prediction = model(offsets_bert, offsets_gcn, bert_embeddings, batched_graph)\n",
    "            if test_predict is None:\n",
    "                test_predict = prediction\n",
    "            else:\n",
    "                test_predict = torch.cat((test_predict, prediction), 0) \n",
    "            loss = loss_func(prediction, labels)\n",
    "            test_loss += loss\n",
    "            # test_loss += loss.detach().item()\n",
    "    \n",
    "    test_loss /= (iter + 1)\n",
    "    print('This fold, the test loss is: ', test_loss)\n",
    "    test_predict_lst.append(test_predict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test Part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_predict_arr = [torch.softmax(pre, -1).clamp(1e-4, 1-1e-4).numpy() for pre in test_predict_lst]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_test_preds = np.mean(test_predict_arr, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import log_loss\n",
    "\n",
    "def extract_target(df):\n",
    "    df[\"Neither\"] = 0\n",
    "    df.loc[~(df['A-coref'] | df['B-coref']), \"Neither\"] = 1\n",
    "    df[\"target\"] = 0\n",
    "    df.loc[df['B-coref'] == 1, \"target\"] = 1\n",
    "    df.loc[df[\"Neither\"] == 1, \"target\"] = 2\n",
    "    return df\n",
    "test_df = extract_target(test_df)\n",
    "log_loss(test_df.target, final_test_preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df_sub = pd.DataFrame(final_test_preds, columns=[\"A\", \"B\", \"NEITHER\"])\n",
    "# df_sub[\"ID\"] = test_df.ID\n",
    "# df_sub = df_sub[['ID',\"A\", \"B\", \"NEITHER\"]]\n",
    "# df_sub.to_csv(\"submission_415_copy3.csv\", index=False)\n",
    "# df_sub.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#df_sub.to_csv(\"submission_414_lg.csv\", index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {
    "height": "246px",
    "width": "337px"
   },
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
